// Copyright 2000-2017 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
package qunar.tc.decompiler.util;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;

public class VBStyleCollection<E, K> extends ArrayList<E> {

    private HashMap<K, Integer> map = new HashMap<>();

    private ArrayList<K> lstKeys = new ArrayList<>();

    public VBStyleCollection() {
        super();
    }

    public VBStyleCollection(int initialCapacity) {
        super(initialCapacity);
        lstKeys = new ArrayList<>(initialCapacity);
        map = new HashMap<>(initialCapacity);
    }

    @Override
    public boolean add(E element) {
        lstKeys.add(null);
        super.add(element);
        return true;
    }

    @Override
    public boolean remove(Object element) {   // TODO: error on void remove(E element)
        throw new RuntimeException("not implemented!");
    }

    @Override
    public boolean addAll(Collection<? extends E> c) {
        for (int i = c.size() - 1; i >= 0; i--) {
            lstKeys.add(null);
        }
        return super.addAll(c);
    }

    public void addAllWithKey(Collection<E> elements, Collection<K> keys) {
        int index = super.size();

        for (K key : keys) {
            map.put(key, index++);
        }

        super.addAll(elements);
        lstKeys.addAll(keys);
    }

    public void addWithKey(E element, K key) {
        map.put(key, super.size());
        super.add(element);
        lstKeys.add(key);
    }

    // TODO: speed up the method
    public E putWithKey(E element, K key) {
        Integer index = map.get(key);
        if (index == null) {
            addWithKey(element, key);
        } else {
            return super.set(index, element);
        }
        return null;
    }

    @Override
    public void add(int index, E element) {
        addToListIndex(index, 1);
        lstKeys.add(index, null);
        super.add(index, element);
    }

    public void addWithKeyAndIndex(int index, E element, K key) {
        addToListIndex(index, 1);
        map.put(key, index);
        super.add(index, element);
        lstKeys.add(index, key);
    }

    public void removeWithKey(K key) {
        int index = map.get(key);
        addToListIndex(index + 1, -1);
        super.remove(index);
        lstKeys.remove(index);
        map.remove(key);
    }

    @Override
    public E remove(int index) {
        addToListIndex(index + 1, -1);
        K obj = lstKeys.get(index);
        if (obj != null) {
            map.remove(obj);
        }
        lstKeys.remove(index);
        return super.remove(index);
    }

    public E getWithKey(K key) {
        Integer index = map.get(key);
        if (index == null) {
            return null;
        }
        return super.get(index);
    }

    public int getIndexByKey(K key) {
        return map.get(key);
    }

    public E getLast() {
        return super.get(super.size() - 1);
    }

    public boolean containsKey(K key) {
        return map.containsKey(key);
    }

    @Override
    public void clear() {
        map.clear();
        lstKeys.clear();
        super.clear();
    }

    @Override
    public VBStyleCollection<E, K> clone() {
        VBStyleCollection<E, K> c = new VBStyleCollection<>();
        c.addAll(new ArrayList<>(this));
        c.setMap(new HashMap<>(map));
        c.setLstKeys(new ArrayList<>(lstKeys));
        return c;
    }

    public void setMap(HashMap<K, Integer> map) {
        this.map = map;
    }

    public K getKey(int index) {
        return lstKeys.get(index);
    }

    public ArrayList<K> getLstKeys() {
        return lstKeys;
    }

    public void setLstKeys(ArrayList<K> lstKeys) {
        this.lstKeys = lstKeys;
    }

    private void addToListIndex(int index, int diff) {
        for (int i = lstKeys.size() - 1; i >= index; i--) {
            K obj = lstKeys.get(i);
            if (obj != null) {
                map.put(obj, i + diff);
            }
        }
    }
}